import numpy as np
from PIL import Image


def calculate_psnr(img1, img2):
    """
    计算两幅图像的峰值信噪比（PSNR）。

    :param img1: 原始图像（NumPy数组）
    :param img2: 需要比较的图像（NumPy数组）
    :return: 峰值信噪比（PSNR）
    """
    # 确保两个图像具有相同的尺寸
    assert img1.shape == img2.shape, 'Input images must have the same dimensions.'

    # 计算均方误差（MSE）
    mse = np.mean((img1 - img2) ** 2)

    # 为了避免返回一个极大的PSNR的值，设置阈值为10^-10，当MSE小于该阈值时，直接返回无穷大
    if mse < 1e-10:
        return float("inf")

    # 峰值信噪比（PSNR）
    # 假设图像的最大像素值为255（对于8位图像）
    MAX_I = 255.0
    PSNR = 20 * np.log10(MAX_I / np.sqrt(mse))

    return PSNR


# 加载图像
def load_image(image_path):
    img = Image.open(image_path).convert('RGB')  # 确保图像是RGB格式
    img = np.array(img, dtype=np.float64)  # 将图像转换为NumPy数组，数据类型为float64
    img = img / 255.0  # 归一化到[0, 1]范围
    return img


# 示例图像路径
image_path1 = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\5.png'
image_path2 = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\ans.png'

# 加载图像
img1 = load_image(image_path1)
img2 = load_image(image_path2)

# 确保两张图像尺寸相同
if img1.shape != img2.shape:
    print('Images have different dimensions. Cannot calculate PSNR.')
else:
    # 计算PSNR
    psnr = calculate_psnr(img1, img2)
    print('PSNR:',psnr)
